package progress

import (
	"io"
	"net/http"
	"path/filepath"

	"webpagedown3/pkg/task"

	"github.com/cheggaaa/pb/v3"
)

// ProgressReader 用于更新进度条
type ProgressReader struct {
	Reader io.Reader
	Bar    *pb.ProgressBar
}

func (pr *ProgressReader) Read(p []byte) (n int, err error) {
	n, err = pr.Reader.Read(p)
	if n > 0 {
		pr.Bar.Add(n)
	}
	return
}

func CreateProgressBar(downloadtask *task.DownloadTask, resp *http.Response) *ProgressReader {
	// 创建进度条
	if downloadtask.Bar == nil {
		if downloadtask.TotalSize > 0 {
			downloadtask.Bar = pb.Full.Start64(downloadtask.TotalSize)
		} else {
			downloadtask.Bar = pb.New(-1)
			downloadtask.Bar.Start()
		}
		downloadtask.Bar.Set("prefix", filepath.Base(downloadtask.URL)+" ")
		downloadtask.Bar.SetCurrent(downloadtask.CurrentSize)
	}

	// 创建代理读取器来更新进度条
	reader := &ProgressReader{
		Reader: resp.Body,
		Bar:    downloadtask.Bar,
	}
	return reader
}
